##########################################################################################

library(data.table)
library(optparse)
library(ArchR)
library(org.Hs.eg.db)
library(GO.db)
library(ggthemes)
library(dplyr)
library(parallel)

##########################################################################################
option_list <- list(
    make_option(c("--gene_list_file"), type = "character"),
    make_option(c("--input_path"), type = "character"),
    make_option(c("--comine_data_all_file"), type = "character"),
    make_option(c("--scriptPath"), type = "character"),
    make_option(c("--out_path"), type = "character") 
)

if(1!=1){
    
    ## 
    gene_list_file <- "~/20231121_singleMuti/results/celltype_plot/trajectory/positive/pct_0.25.list"

    ## 所有的细胞的
    input_path <- "~/20231121_singleMuti/results/celltype_plot/tf_regulators_all_germ"

    ## 所有的细胞的
    comine_data_all_file <- "~/20231121_singleMuti/results/qc_atac_v3/germ/testis_combined_peak.combineRNA.qc.Rdata"

    ## 既往研究整理的代码
    scriptPath <- "~/20231121_singleMuti/scripts/scScalpChromatin"

    ## 输出
    out_path <- "~/20231121_singleMuti/results/celltype_plot/tf_regulators_all_germ_pathway"

}

###########################################################################################
parseobj <- OptionParser(option_list=option_list, usage = "usage: Rscript %prog [options]")
opt <- parse_args(parseobj)
print(opt)

gene_list_file <- opt$gene_list_file
input_path <- opt$input_path
comine_data_all_file <- opt$comine_data_all_file
scriptPath <- opt$scriptPath
out_path <- opt$out_path

regPlotDir <- out_path
dir.create(regPlotDir, showWarnings = FALSE, recursive = TRUE)

#######################################################################################
## 导入数据

gene_list <- data.frame(fread(gene_list_file , header = F))$V1
b <- load(comine_data_all_file)
cpu <- 20

###########################################################################################
## 已发表文献写好的脚本
source(paste0(scriptPath, "/plotting_config.R"))
source(paste0(scriptPath, "/misc_helpers.R"))
source(paste0(scriptPath, "/matrix_helpers.R"))
source(paste0(scriptPath, "/archr_helpers.R"))
source(paste0(scriptPath, "/GO_wrappers.R"))

###########################################################################################
## 获取背景基因
if(exists("testis_combined_peak_combineRNA")){
  atac_proj <- testis_combined_peak_combineRNA
}

# GeneIntegration Matrix: (rows gene names x cols cell names)
GIMatrix <- getMatrixFromProject(atac_proj, useMatrix="GeneExpressionMatrix")
GImat <- assays(GIMatrix)$GeneExpressionMatrix
rownames(GImat) <- rowData(GIMatrix)$name
# Remove unexpressed genes
GImat <- as(GImat[Matrix::rowSums(GImat) > 0,], "sparseMatrix") 

# Already filtered to only expressed genes
allGenes <- rownames(GImat) %>% sort() 

###########################################################################################

mclapply(gene_list,function(gene){

  input_file <- paste0( input_path , "/" , gene , "_LS.tsv" )
  dat_ls <- data.frame(fread(input_file))

  ###########################################################################################

  motif_short <- gene
  maxPval <- 5

  ##########################################################################################
  # plot all TF regulators

  # Store results for each TF
  res_list <- list()

  ## 显著调控基因的阈值
  plot_df <- dat_ls
  LS_window <- quantile(plot_df$LS, 0.8)
  corr_window <- 0.25
  mLog10pval_window <- quantile(plot_df$mLog10pval, 0.8)

  ## 考虑两个维度
  pos_top_genes <- plot_df[plot_df$LS > LS_window & plot_df$Correlation > corr_window,]$symbol
  ## 考虑三个维度
  pos_top_enrich_genes <- plot_df[plot_df$LS > LS_window & plot_df$Correlation > corr_window & plot_df$mLog10pval > mLog10pval_window,]$symbol

  ##########################################################################################
  ## 画图
  plot_df$putative_targets_cor_ls <- ifelse( plot_df$symbol %in% pos_top_genes , "YES" , "NO" )
  plot_df$putative_targets_cor_ls_enrich <- ifelse( plot_df$symbol %in% pos_top_enrich_genes , "YES" , "NO" )

  for( uselabel in c("putative_targets_cor_ls" , "putative_targets_cor_ls_enrich") ){

    plot_df$toLabel <- plot_df[[uselabel]]
    use_gene <- subset(plot_df , toLabel=="YES")$symbol

    ## plot
    topN <- 30
    plot_df$toLabel <- "NO"
    plot_df <- plot_df[order(plot_df$LS , plot_df$Correlation , decreasing=TRUE),]

    if( uselabel == "putative_targets_cor_ls_enrich" ){
      plot_df <- plot_df[order(plot_df$LS , plot_df$Correlation , plot_df$mLog10pval, decreasing=TRUE),]
    }

    plot_df$toLabel[1:topN] <- "YES"
    plot_df$toLabel <- ifelse( plot_df$symbol %in% use_gene , plot_df$toLabel , "NO" )

    p <- (
      ggplot(plot_df, aes(x=Correlation, y=LS, color=mLog10pval)) 
        #+ geom_point(size = 2)
        + ggrastr::geom_point_rast(size=2)
        + ggrepel::geom_text_repel(
            data=plot_df[plot_df$toLabel=="YES",], aes(x=Correlation, y=LS, label=symbol), 
            #data = plot_df, aes(x=Correlation, y=LS, label=symbol), #(don't do this, or the file will still be huge...)
            size=2,
            point.padding=0, # additional pading around each point
            box.padding=0.5,
            min.segment.length=0, # draw all line segments
            max.overlaps=Inf, # draw all labels
            #nudge_x = 2,
            color="black"
        ) 
        + geom_vline(xintercept=corr_window, lty="dashed", color="red")
        + geom_hline(yintercept=LS_window, lty="dashed", color="red")
        + theme_BOR(border=FALSE)
        + theme(panel.grid.major=element_blank(), 
                panel.grid.minor= element_blank(), 
                plot.margin = unit(c(0.25,1,0.25,1), "cm"), 
                aspect.ratio=1.0,
                #legend.position = "none", # Remove legend
                axis.text.x = element_text(angle=90, hjust=1))
        + ylab("Linkage Score") 
        + xlab("Motif Correlation to Gene") 
        + scale_color_gradientn(colors=cmaps_BOR$zissou, limits=c(0, maxPval))
        + scale_y_continuous(expand = expansion(mult=c(0,0.05)))
        + scale_x_continuous(limits = c(-0.85, 0.955)) # Force plot limits
        + ggtitle(sprintf("%s putative targets", motif_short))
    )

    ## 通路富集
    if( length(use_gene) > 0 ){
      upGO <- rbind(
        calcTopGo(allGenes, interestingGenes=use_gene, nodeSize=5, ontology="BP") 
      )

      upGO <- upGO[order(as.numeric(upGO$pvalue), decreasing=FALSE),]

      ## 构造GO对象
      geneList <- factor(as.integer(allGenes %in% use_gene))
      names(geneList) <- allGenes
      GOdata <- suppressMessages(new(
        "topGOdata",
        ontology = "BP",
        allGenes = geneList,
        annot = annFUN.org, mapping = "org.Hs.eg.db", ID = "symbol",
        nodeSize = 5
      ))

      ## 提取每个通路中感兴趣的基因
      gene_in_pathway <- sapply(upGO$GO.ID, function(x)
        {
          genes <- genesInTerm(GOdata, x)
          # myGenes is the queried gene list
          paste0(genes[[1]][genes[[1]] %in% use_gene] , collapse = "," )
        })

      upGO$gene_in_pathway <- gene_in_pathway

      ## 可视化
      if(min(upGO$pvalue) < 1){
        upGO <- upGO[order(as.numeric(upGO$pvalue), decreasing=FALSE),]
        up_go_plot <- topGObarPlot(upGO, cmap=cmaps_BOR$comet, nterms=10, border_color="black", 
          barwidth=0.9, title=sprintf("%s putative targets (%s genes)", motif_short, length(use_gene)), enrichLimits=c(0, 6))
      }else{
        up_go_plot <- ggplot()
      }

      ## 输出通路
      out_file <- paste0(regPlotDir, "/" , motif_short , "_GO." , uselabel , ".tsv")
      write.table( upGO , out_file , row.names = F , quote = F , sep ="\t" ) 

    }

    ## 输出图片
    out_file <- paste0(regPlotDir, "/" , motif_short , "." , uselabel , ".pdf")
    pdf(out_file, width=8, height=6)
    print(p)
    if( length(pos_top_genes) > 0 ){
      print(up_go_plot)
    }
    dev.off()

    ## 输出
    use_col <- !colnames(plot_df) %in% c("putative_targets" , "label_target" , "toLabel")
    out_file <- paste0(regPlotDir, sprintf("/%s_LS.tsv", motif_short))
    write.table(plot_df[,use_col] , out_file , row.names = F , sep = "\t" , quote = F)
  }

},mc.cores=cpu)